"""
Interface definitions for backbone networks in continual learning.

Backbone networks are feature extractors that form the foundation of
continual learning models. This module defines interfaces ensuring
consistent behavior across different backbone implementations.
"""

from abc import abstractmethod

import torch
from torch import nn
from continuallearning.interfaces.core.component import ComponentInterface
from continuallearning.interfaces.types import BackboneOutput


class BackboneInterface(nn.Module, ComponentInterface):
    """
    Interface for model backbone implementations.

    Backbones are responsible for feature extraction from input data.
    They serve as the foundation for continual learning models and can be
    adapted through various approaches like parameter-efficient fine-tuning.
    """

    @abstractmethod
    def forward(self, x: torch.Tensor, **kwargs) -> BackboneOutput:
        """
        Extract features from input data.

        Args:
            x: Input tensor typically of shape [B, C, H, W] or [B, seq_len, hidden_size]
            **kwargs: Additional arguments for flexibility

        Returns:
            Union[Dict[str, Any], BackboneOutput]: Output containing features and optional
                                            hidden states or attention weights
        """
        pass

    @abstractmethod
    def get_feature_dims(self) -> int:
        """
        Get the output dimension of the backbone features.

        Returns:
            int: Feature dimension
        """
        pass

    def extract_features(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        """
        Extract features from input without additional processing.

        This convenience method calls forward and returns just the features.

        Args:
            x: Input tensor
            **kwargs: Additional arguments for the backbone

        Returns:
            torch.Tensor: Feature tensor representation
        """
        output = self.forward(x, **kwargs)

        if isinstance(output, dict):
            if "features" in output:
                return output["features"]
            elif "last_hidden_state" in output:
                # Use CLS token from transformer models
                return output["last_hidden_state"][:, 0]
            else:
                # Use the first tensor found
                for key, value in output.items():
                    if isinstance(value, torch.Tensor):
                        return value
        elif hasattr(output, "features") and output.features is not None:
            return output.features
        elif isinstance(output, torch.Tensor):
            return output

        raise ValueError(f"Unexpected output type: {type(output)}")
